
import pygame
import numpy as np
import sounddevice as sd
from math import pi, tau
import random
deltakey = 2 ** (1 / 12)
a1 = 440
SAMPLE_RATE = 44100
out = sd.default.device[1]
pygame.init()
def get_hz(n):
    return a1 * (deltakey ** n)

def to_short(arr):
    assert arr.dtype == np.float64
    r = (np.clip(arr, -1, 1) * 32767).astype(np.int16)
    return r

class Control:
    pass


class Window:
    def __init__(self, size):
        self.size = size
        self.size_half = (size[0] // 2, size[1] // 2)
        self.screen = pygame.display.set_mode(size)

        self.draw2d_list = []

        self.key_down_callback = []
        self.key_up_callback = []
        self.mouse_down_callback = []
        self.mouse_up_callback = []
        self.mouse_move_callback = []    # 一般只使用第一个元素
        self.update_callback = []

    def __rrshift__(self, obj):
        if isinstance(obj, Control):
            if hasattr(obj, "draw"):
                self.draw2d_list.append(obj)
            if hasattr(obj, "reg_event"):     # 并不能确定要注册多少种事件，并且事件可以临时删除
                obj.reg_event(self)
        else:
            raise Exception(f" >> window fail, {obj}")

    def dispatch(self, e):
        if e.type == pygame.MOUSEBUTTONDOWN:
            for callback in self.mouse_down_callback:
                r = callback(self, e)
                if r == True:
                    break
        elif e.type == pygame.MOUSEBUTTONUP:
            for callback in self.mouse_up_callback:
                r = callback(self, e)
                if r == True:
                    break
        elif e.type == pygame.MOUSEMOTION:
            for callback in self.mouse_move_callback:
                callback(self, e)
        elif e.type == pygame.KEYDOWN:
            for callback in self.key_down_callback:
                r = callback(self, e)
                if r == True:
                    break
        elif e.type == pygame.KEYUP:
            for callback in self.key_up_callback:
                r = callback(self, e)
                if r == True:
                    break

    def render(self):
        self.screen.fill((0, 0, 0))
        for control in self.draw2d_list:
            control.draw(self)

    def update(self, dt):
        for callback in self.update_callback:
            callback(self, dt / 1000)

class GuitaKeyFade(Control):
    def __init__(self, app):
        self.d_down = {}
        self.d_up = {}
        self.chord = [None, None, None, None, None, None, None]

        self.app = app  # 用于回调加入音符

        def add_down(key):
            def g(f):
                self.d_down[key] = f
            return g
        def add_up(key):
            def g(f):
                self.d_up[key] = f
            return g
        # 右手
        @add_down(pygame.K_KP0)
        def _(self):
            if self.chord[0]:
                self.sound(1, self.chord[0])
        @add_down(pygame.K_KP1)
        def _(self):
            if self.chord[0]:
                self.sound(2, self.chord[1])
        @add_down(pygame.K_KP4)
        def _(self):
            if self.chord[0]:
                self.sound(3, self.chord[2])
        @add_down(pygame.K_KP7)
        def _(self):
            if self.chord[0]:
                self.sound(4, self.chord[3])
        @add_down(pygame.K_KP8)
        def _(self):
            if self.chord[0]:
                self.sound(5, self.chord[4])
        @add_down(pygame.K_KP9)
        def _(self):
            if self.chord[0]:
                self.sound(6, self.chord[5])
        @add_down(pygame.K_KP_PLUS)
        def _(self):
            if self.chord[0]:
                self.sound(7, self.chord[6])

        # 左手
        
        def chord3(x, tag):
            return [x - 12, x - 8, x - 5, x, x + 4, x + 7, x + 12], tag
        
        def chord3m(x, tag):
            return [x - 12, x - 9, x - 5, x, x + 3, x + 7, x + 12], tag
        
        d = {
        pygame.K_1 : chord3(-9, "C"),
        pygame.K_2 : chord3(-7, "D"),
        pygame.K_3 : chord3(-5, "E"),
        pygame.K_4 : chord3(-4, "F"),
        pygame.K_5 : chord3(-2, "G"),
        pygame.K_6 : chord3(0, "A"),
        pygame.K_7 : chord3(2, "B"),
        pygame.K_q : chord3m(-9, "Cm"),
        pygame.K_w : chord3m(-7, "Dm"),
        pygame.K_e: chord3m(-5, "Em"),
        pygame.K_r : chord3m(-4, "Fm"),
        pygame.K_t : chord3m(-2, "Gm"),
        pygame.K_y : chord3m(0, "Am"),
        pygame.K_u : chord3m(2, "Bm"),
        }
        
        for key, chord in d.items():
            @add_down(key)
            def _(self, chord=chord):
                print(chord[1])
                self.chord = chord[0]
            @add_up(key)
            def _(self, chord=chord):
                if self.chord == chord[0]:
                    self.chord = [None, None, None, None, None, None, None]
        

    def sound(self, line, num):
        note = self.app.note_type(get_hz(num))
        self.app.get_line(line).load(note)

    # ~ def sound_end(self, line):
        # ~ self.app.get_line(line).stop()

    # ~ def sound_change(self, line, num):
        # ~ note = self.app.note_type(get_hz(num))
        # ~ self.app.get_line(line).change(note)

    def key_down(self, window, e):
        callback = self.d_down.get(e.key, lambda self:None)
        callback(self)

    def key_up(self, window, e):
        callback = self.d_up.get(e.key, lambda self:None)
        callback(self)

    def reg_event(self, window):
        window.key_down_callback.append(self.key_down)
        window.key_up_callback.append(self.key_up)


class MusicLine:
    def __init__(self, stream):
        self.stream = stream
        self.note = None
        self.note_start_time = None

    def load(self, note):
        # ~ print("load")
        self.note = note
        self.note_start_time = self.stream.time

    def change(self, note):
        # ~ print("change")
        self.note = note

    def stop(self):
        # ~ print("stop")
        self.note = None
        self.note_start_time = None

    def pop(self, t, length):
        note_start_time = self.note_start_time       # 此函数会在其他线程调用！
        note = self.note
        if note is None or note_start_time is None:
            return np.zeros(length)
        else:
            tl = np.arange(length) / SAMPLE_RATE
            tl = tl + (t - note_start_time)
            b, arr = note.f(tl)
            if b == False:
                self.stop()
            return arr


class Note:
    def f(self, tl)    -> (bool, np.ndarray):         # 返回False删除该音
        raise NotImplemented()

def sign2(arr):
    l2 = np.zeros(len(arr))
    b = arr > l2
    return b.astype(np.float64)

class SineUp(Note):
    def __init__(self, f0, k):
        self.f0 = f0
        self.k = k
        self.delta = 2 ** (1 / 12)

    def f(self, tl):
        l = np.sin(tl * tau * (self.f0 * (self.delta ** (tl * self.k))))
        l = sign2(tl) * l
        return True, l

#su = SineUp(440, 12)

class Sine(Note):
    def __init__(self, f0):
        self.f0 = f0

    def f(self, tl):
        l = np.sin(tl * tau * self.f0)
        l = sign2(tl) * l
        return True, l

class Sine2(Note):
    def __init__(self, f0):
        self.f0 = f0

    def f(self, tl):
        l = np.sin(tl * tau * self.f0)
        l += 0.5 * np.sin(tl * tau * self.f0 * 2)
        l += 0.3 * np.sin(tl * tau * self.f0 * 3)
        l += 0.1 * np.sin(tl * tau * self.f0 * 4)
        l = sign2(tl) * l
        return True, l

def fade(tl, k, b):     # b 秒之后归零
    l = np.exp(-tl * k) - np.exp(-k * b)
    l = sign2(tl) * l
    return l

class SineFade(Note):
    def __init__(self, f0):
        self.f0 = f0

    def f(self, tl):
        l = np.sin(tl * tau * self.f0)
        l = fade(tl, 0.5, 5) * l
        return tl[-1] < 5, l

class Sine2Fade(Note):
    def __init__(self, f0):
        self.f0 = f0

    def f(self, tl):
        l = np.sin(tl * tau * self.f0)
        l += 0.5 * np.sin(tl * tau * self.f0 * 2)
        l += 0.4 * np.sin(tl * tau * self.f0 * 3)
        l += 0.3 * np.sin(tl * tau * self.f0 * 4)
        l += 0.2 * np.sin(tl * tau * self.f0 * 5)
        l += 0.1 * np.sin(tl * tau * self.f0 * 6)
        l += 0.1 * np.sin(tl * tau * self.f0 * 7)
        l = fade(tl, 0.5, 5) * l
        return tl[-1] < 5, l

base = get_hz(-24)

feature_energy = {
1: 0.75, 
2: 3.5, 
3: 9, 
4: 0, 
5: 0, 
6: 0, 
7: 1.66, 
8: 0.8, 
9: 0.3, 
10: 0, 
11: 0, 
12: 0, 
13: 0, 
14: 1.66, 
15: 3.08, 
16: 0, 
17: 0, 
18: 6.3, 
19: 32, 
20: 14, 
21: 2.3,
22: 0,
23: 0,
24: 0,
25: 2.1,
26: 3.6,
27: 1.8,
28: 0,
29: 0.7,
30: 1.1,
31: 2.6,
32: 0.8,


}


class Sine4Fade(Note):          # energy k
    f_list = np.array(range(1, 33), dtype=np.float64)
    def __init__(self, f0):
        self.f0 = f0
        k = f0 / base
        print("k=",k)
        l = [*feature_energy.values()]
        for i in range(len(l)):
            l[i] =( l[i]  / (i +1) / (i + 1)) ** 0.5
            
        print("!", [f"{n:.3f}" for n in l])
        # ~ temp = sum(l)
        # ~ for i in range(len(l)):
            # ~ l[i] = l[i]  / temp
        # ~ print([f"{n:.3f}" for n in l])
            
        self.fe_list = np.array(l, dtype=np.float64)

    def f(self, tl):
        l = np.sin(tau * self.f0 * np.outer(tl, self.f_list))
        l = l.dot(self.fe_list)
        l = fade(tl, 0.5, 5) * l
        return tl[-1] < 5, l
        
        
class Sine3Fade(Note):      # base k
    f_list = np.array(range(1, 41), dtype=np.float64)
    # ~ fe_list = np.array([*feature.values()], dtype=np.float64)
    def __init__(self, f0):
        self.f0 = f0

    def f(self, tl):
        l = np.sin(tau * self.f0 * np.outer(tl, self.f_list))
        l = l.dot(self.fe_list)
        l = fade(tl, 0.5, 5) * l
        return tl[-1] < 5, l

class App:
    def __init__(self, guitakey_type, note_type):
        def callback(outdata, frames, time, status):
            sound1 = self.ml1.pop(time.outputBufferDacTime, frames)
            sound2 = self.ml2.pop(time.outputBufferDacTime, frames)
            sound3 = self.ml3.pop(time.outputBufferDacTime, frames)
            sound4 = self.ml4.pop(time.outputBufferDacTime, frames)
            sound5 = self.ml5.pop(time.outputBufferDacTime, frames)
            sound6 = self.ml6.pop(time.outputBufferDacTime, frames)
            sound7 = self.ml7.pop(time.outputBufferDacTime, frames)
            sound = (sound1 + sound2 + sound3 + sound4 + sound5 + sound6 + sound7) / 4
            sound = to_short(sound)
            multiChannel = np.column_stack((sound, sound))
            outdata[:] = multiChannel
        self.s1 = sd.OutputStream(samplerate=SAMPLE_RATE, device= out, channels=2, dtype=np.int16, blocksize=2048, callback=callback)

        self.ml1 = MusicLine(self.s1)
        self.ml2 = MusicLine(self.s1)
        self.ml3 = MusicLine(self.s1)
        self.ml4 = MusicLine(self.s1)
        self.ml5 = MusicLine(self.s1)
        self.ml6 = MusicLine(self.s1)
        self.ml7 = MusicLine(self.s1)



        self.window = Window((1000, 800))
        self.note_type = note_type
        self.guitakey = guitakey_type(self)
        self.guitakey >> self.window

        self.clock = pygame.time.Clock()

        self.s1.start()

    def get_line(self, n):
        return [self.ml1, self.ml2, self.ml3, self.ml4, self.ml5, self.ml6, self.ml7][n-1]

    def run(self):
        while True:
            for e in pygame.event.get():
                if e.type == pygame.QUIT:
                    exit()
                self.window.dispatch(e)
            self.window.render()
            pygame.display.flip()
            dt = self.clock.tick(60)
            self.window.update(dt)

# ~ app = App(GuitaKeyFade, Sine2Fade)
# ~ app = App(GuitaKeyFade, Sine2Fade)
# ~ app = App(GuitaKeyFade, Sine3Fade)
app = App(GuitaKeyFade, Sine4Fade)
app.run()

